import sglang as sgl
import re
import json
import random
from .cann_data import *
from kernel_gen.logger import init_logger
logger = init_logger(__name__)
def get_prompt(chapter_id="1.1"):
    if chapter_id.find("16.") is not None:
        prompt = """你是一位精通文本分析的AI助手。我将提供一段文档章节，这个章节是ascendc的api，你的任务是深入理解这段内容，并从中提取出10个不重复的问题。
    请遵循以下要求：
    1. 仔细阅读和理解所提供的文档内容，把握其核心观点、关键概念和逻辑结构。
    2. 基于文档内容生成10个有价值的问题，这些问题应该：
    - 覆盖文档的主要内容和细节
    - 避免表述模糊或过于宽泛的问题
    - 确保每个问题都基于文档中的实际内容
    - 确保每个问题都能从文档中直接找到答案
    3. 所有问题必须使用以下格式严格输出：
    <q>问题1</q>
    <q>问题2</q>
    <q>问题3</q>
    ...
    4. 确保所有问题之间不存在重复或高度相似的问题，每个问题应该针对文档的不同方面或角度。
    6. 请直接生成符合要求的10个问题，不需要其他解释或说明。
    7. 提问内容不要涉及CUDA, OPENCV等其他框架或硬件体系。
    8. 多提问公式、api用法、使用场景、输入输出范式、示例代码和使用限制等问题。
    8. 不需要提问接口实现细节的问题。
    """
    else:
        prompt = f"""你是一位精通文本分析的AI助手。我将提供一段文档章节，你的任务是深入理解这段内容，并从中提取出10个不重复的问题。
    请遵循以下要求：
    1. 仔细阅读和理解所提供的文档内容，把握其核心观点、关键概念和逻辑结构。
    2. 基于文档内容生成10个有价值的问题，这些问题应该：
    - 覆盖文档的主要内容和细节
    - 包含不同认知层次的问题（如理解、分析、应用、评价等）
    - 避免表述模糊或过于宽泛的问题
    - 确保每个问题都基于文档中的实际内容
    - 多提问和代码相关的问题
    - 多提问公式和代码进行提问
    3. 所有问题必须使用以下格式严格输出：
    <q>问题1</q>
    <q>问题2</q>
    <q>问题3</q>
    ...
    4. 确保所有问题之间不存在重复或高度相似的问题，每个问题应该针对文档的不同方面或角度。
    5. 如果文档内容有限，也请尽力从不同角度提出问题，包括推理性、假设性问题，但这些问题仍应与文档内容相关。
    6. 请直接生成符合要求的10个问题，不需要其他解释或说明。
    7. 提问内容不要涉及CUDA, OPENCV等其他框架或硬件体系。
    8. 如果章节内容是api介绍，请你多提问api用法、使用场景、输入输出、使用限制等问题。
    """        
    return prompt

def slice_questions(questions_str:str):
    matches = re.findall(r"<q>(.*?)</q>", questions_str, re.DOTALL)
    ret = []
    if len(matches) == 0:
        ret.append(questions_str)
    else:
        for match in matches:
            ret.append(match)
    return ret

@sgl.function
def ask_question(s,section_test,question_num = 10):
    # s += (
    #     f"下面是昇腾文档的{section_id}节的所有代码段，请你仔细阅读，并提出对文章内容进行提问。提的问题需要有深度和广度，注意你的问题一定是能在文档中找到答案或者可以通过联系文档中内容推导出答案，你的提问只需要包含问题本身，保证一次生成一个问题，并且不要有任何其他的内容，最后用?或者？作为结尾。"
    #     f" {section.to_prompt()}\n\n"
    # )
    
    s += sgl.system(get_prompt())
    s += sgl.user(f"下面是昇腾文档的部分内容：{section_test}，请按照要求进行提问。")
    forks = s.fork(question_num)
    # for i, f in enumerate(forks):
    # forks += f"现在请你提问第 {i+1} 个问题，不要和之前重复:\n"
    forks += sgl.assistant(sgl.gen(f"question", max_tokens=4096,temperature=1.2,))
    forks.join()

@sgl.function
def ask_question_multi_round(s,section_test,max_round = 20, fork_num = 1):
    s += sgl.system(get_prompt())
    s += sgl.user(f"下面是昇腾文档的部分内容：{section_test},请按照要求进行提问。")
    forks = s.fork(fork_num)
    for round in range(max_round):
        if round > 0:
            forks += sgl.user("请你继续按照要求进行提问，不要与之前问题重复。")            
        forks += sgl.assistant(sgl.gen(f"question_{round}", max_tokens=4096,temperature=1.2))
    forks.join()

@sgl.function
def ask_question_multi_round(s,section_test,max_round = 20, fork_num = 1):
    s += sgl.system(get_prompt())
    s += sgl.user(f"下面是昇腾文档的部分内容：{section_test},请按照要求进行提问。")
    for round in range(max_round):
        if round > 0:
            s += sgl.user("请你继续按照要求进行提问，不要与之前问题重复。")            
        s += sgl.assistant(sgl.gen(f"question_{round}", max_tokens=4096,temperature=1.2))


@sgl.function()
def answer_question(s,section_id,section,question:str):
    s += sgl.user(f"下面是文档的第{section_id}节的所有内容：{section.to_prompt()} ， 请你标注该问题涉及到的章节，无论单个章节还是多个章节都请以数组形式列举，并保证json可以正常读取。问题如下：{question}")
    s += sgl.assistant(sgl.gen(f"answer_in_sections", max_tokens=128,temperature=1.2,))

def ask_question_pipeline_multi_round(root_section:Section,max_round = 20, fork_num = 1, file_name = "ds_questions_list.json"):
    section_text = ""
    section_idx = []
    logger.info(f"Start asking question for {root_section.id},conv round= {max_round}, fork_num={fork_num},file_name={file_name}")
    for key,current_section in root_section.all_item_with_depth.items():
        section_text += current_section.to_prompt() +"\n"
        section_idx.append(key)
    try:
        status = ask_question_multi_round.run(section_test=section_text,max_round = max_round, fork_num = fork_num,stream=True)
        status.sync()
    except Exception as e:
        logger.error(f"error message:{e}")
    logger.info(status.messages())
    all_question_list = []
    for round in range(max_round):
        # for i in range(fork_num):
        question = status[f"question_{round}"]
        q_list = slice_questions(question)
        all_question_list.extend(q_list)
    logger.info(f"End asking question for {root_section.id}, Start to Saving json File, file name= {file_name}")
    router = save_questions_to_json(file_name)
    for idx,current_q in enumerate(all_question_list):
        logger.info(f"question {idx}:{current_q}")
        router["add_new_question"](current_q, max_range_questions=section_idx)
    router["save_to_file"]()
    logger.info(f"End saving json File for {root_section.id}, file name= {file_name}")
    


def ask_question_pipeline(section_text,question_num = 10, file_name = "ds_questions_list.json",need_answer=False):
    status = ask_question.run(section_test=section_text,question_num=question_num)
    all_question_list = []
    for i in range(question_num):
        question = status["question"][i]
        q_list = slice_questions(question)
        all_question_list.extend(q_list)
    print(status.text())
    router = save_questions_to_json(file_name)
    for idx,current_q in enumerate(all_question_list):
        print(f"question {idx}:{current_q}")
        router["add_new_question"](current_q,None)
    router["save_to_file"]()

@sgl.function
def ask_question_for_codes(s,section_ids,sections,question_num = 10):
    if not isinstance(section_ids,list):
        section_ids = [section_ids]
    if not isinstance(sections,list):
        sections = [sections]
    sections_codes = ""
    for section in sections:
        sections_codes+=section.get_all_code_snippets()
    s += (
        f"下面是昇腾文档的{section_ids}节的所有代码段，请你仔细阅读，并提出对这段代码内容进行提问。注意，提问的内容需要包含代码逻辑，接口用法，存储结构，执行效率等多方面内容，保证一次生成一个问题，并且不要有任何其他的内容，最后用?或者？作为结尾。"
        f" {sections_codes}\n\n"
    )
    
    forks = s.fork(question_num)
    # for i, f in enumerate(forks):
    # forks += f"现在请你提问第 {i+1} 个问题，不要和之前重复:\n"
    forks += sgl.gen(f"question",stop=["\n\n"], max_tokens=256,temperature=1.2,)

    forks += "请你标注该问题涉及到的章节，无论单个章节还是多个章节都请以数组形式列举，并保证json可以正常读取。" + sgl.gen(f"answer_in_sections", max_tokens=256, stop="\n",temperature=1.2,)
    forks.join()


def save_questions_to_json(filename="question_jsons.json"):
    idx = 0
    output_dict = {}
    json_file_name = filename
    def add_new_question(q:str,a = None,max_range_questions = None):
        nonlocal output_dict,idx,json_file_name
        if max_range_questions is None:
            dict_value = q
        else:
            dict_value = {
                "question":q,
                "max_range_questions":max_range_questions
            }
            if a is not None:
                dict_value["llm_anser"] = a
        output_dict[idx] = dict_value
        idx += 1

    def save_to_file():
        nonlocal output_dict,idx,json_file_name
        with open(json_file_name,"w") as f:
            json.dump(output_dict,f,ensure_ascii=False, indent=4)

    return {"add_new_question":add_new_question,"save_to_file":save_to_file}

def ask_question_for_one_chapter_all_section(chapter:dict):
    '''
    针对每一个section提问
    '''
    question_num=1
    router = save_questions_to_json("ds_qa.json")
    for key in chapter.keys():
        value = chapter[key]
        for k in value.children.keys():
            s = value.children[k]
            state =  ask_question.run(section_id=k,section=s,question_num=question_num)
            print(f"running {k} {s}:")
            print(state.text())
            for i in range(question_num):
                print(f"question {i}:",state["question"][i])
                print(f"答案所在的章节:",state["answer_in_sections"][i])
            router["add_new_question"](state,question_num,k)
            print("=="*50)
        break
    router["save_to_file"]()

def ask_question_about_codes_for_one_chapter_all_section(chapter:dict):
    '''
    针对每一个section的代码进行提问
    '''
    question_num=5
    router = save_questions_to_json("questions_for_codes.json")
    for key in chapter.keys():
        value = chapter[key]
        for k in value.children.keys():
            s = value.children[k]
            state =  ask_question_for_codes.run(sections = s,question_num=question_num)
            print(f"running {k} {s}:",flush=True)
            print(state.text())
            for i in range(question_num):
                print(f"question {i}:",state["question"][i],flush=True)
                print(f"答案所在的章节:",state["answer_in_sections"][i],flush=True)
            router["add_new_question"](state,question_num,k)
            print("=="*50,flush=True)
    router["save_to_file"]()

def ask_question_for_one_chapter_pipeline(chapter:Book):
    '''
    针对每一个section的代码进行提问
    '''
    question_num=20
    for key in chapter.children.keys():
        value = chapter[key]
        sections_text = ""
        section_list = []
        for key,current_section in value.all_item_with_depth.items():
            sections_text += current_section.to_prompt() +"\n"
        ask_question_pipeline(section_text = sections_text,file_name=f"qwq_ask_chapter_{key}_question.json",question_num=question_num)
        print("=="*50)

def ask_question_for_one_chapter_all_section_pipeline(chapter:dict):
    '''
    针对每一个section进行提问
    '''
    question_num=20
    for key in chapter.keys():
        value = chapter[key]
        for k in value.children.keys():
            s = value.children[k]
            sections_text = ""
            section_list = []
            for key,current_section in s.all_item_with_depth.items():
                sections_text += current_section.to_prompt() +"\n"
            ask_question_pipeline(section_text = sections_text,file_name=f"ds_ask_chapter_{k}_question.json",question_num=question_num)
            print("=="*50)

def ask_question_for_one_chapter_all_section_pipeline_random(chapter_0:dict,chapter_1:dict):
    '''
    多个section进行提问
    '''
    question_num=100
    sections_text = ""
    for key,current_section in chapter_0.all_item_with_depth.items():
        sections_text += current_section.to_prompt() +"\n\n"
    for key,current_section in chapter_1.all_item_with_depth.items():
        sections_text += current_section.to_prompt() +"\n\n"
    # print(sections_text)
    ask_question_pipeline(section_text = sections_text,file_name=f"ds_ask_multi_chapter_questions.json",question_num=question_num)
    print("=="*50)

def ask_question_for_combined_section(all_chapters):
    '''
    所有section任选n个提问，未完工
    '''
    print(all_chapters)
    all_section = []
    for key in all_chapters.keys():
        value = all_chapters[key]
        for k in value.children.keys():
            s = value.children[k]
            all_section.append((k,s))
    print(all_section)
    n = 3  # 选择的元素个数
    selected_elements = random.sample(all_section, n)

# TODO 针对代码提问，将代码段放回原文提问，针对每一个Chapter提问

    

    

